#include "sm3.h"

#include <stdio.h>
#include <string.h>
#include <stdint.h>

#include <stdlib.h>

// 定义一些常量，符合SM3算法标准
#define SM3_BLOCK_SIZE 64
#define SM3_DIGEST_SIZE 32

// 初始向量，SM3算法规定的初始哈希值
const uint32_t IV[8] = {
    0x7380166F, 0x4914B2B9, 0x172442D7, 0xDA8A0600,
    0xA96F30BC, 0x163138AA, 0xE38DEE4D, 0xB0FB0E4E
};

const uint32_t T_j_1 = 0x79CC4519;
const uint32_t T_j_2 = 0x7A879D8A;

// 循环左移操作函数
uint32_t ROTL(uint32_t x, int n) {
    return ((x << n) | (x >> (32 - n)));
}

// 用于SM3算法中的布尔函数FF
uint32_t FF(uint32_t X, uint32_t Y, uint32_t Z, int j) {
    if (j >= 0 && j <= 15) {
        return X ^ Y ^ Z;
    } else {
        return (X & Y) | (X & Z) | (Y & Z);
    }
}

// 用于SM3算法中的布尔函数GG
uint32_t GG(uint32_t X, uint32_t Y, uint32_t Z, int j) {
    if (j >= 0 && j <= 15) {
        return X ^ Y ^ Z;
    } else {
        return (X & Y) | (~X & Z);
    }
}

// 置换函数P0
uint32_t P0(uint32_t X) {
    return X ^ ROTL(X, 9) ^ ROTL(X, 17);
}

// 置换函数P1
uint32_t P1(uint32_t X) {
    return X ^ ROTL(X, 15) ^ ROTL(X, 23);
}

// 消息填充函数，使消息长度符合SM3算法要求，适配32位系统
// void padding(uint8_t *message, size_t message_size, uint8_t *padded_message, size_t *padded_size) {
//     size_t k = (SM3_BLOCK_SIZE - (message_size + 1 + 8) % SM3_BLOCK_SIZE) % SM3_BLOCK_SIZE;
//     *padded_size = message_size + 1 + k + 8;
//     memset(padded_message, 0, *padded_size);
//     memcpy(padded_message, message, message_size);
//     padded_message[message_size] = 0x80;
//     uint32_t bit_length = (uint32_t)(message_size * 8);  // 强制转换为32位无符号整数，避免溢出等问题
//     for (int i = 0; i < 8; i++) {
//         padded_message[*padded_size - 8 + i] = (uint8_t)(bit_length >> (24 - 8 * i));  // 修改位运算，适配32位整数的字节序提取
//     }
// }


uint8_t* padding(const uint8_t *message, size_t message_size, size_t *padded_size)
{
    // char *padded = 0;
    uint8_t *padded = NULL;;
    int r = message_size % SM3_BLOCK_SIZE;
    int k = 0; 
    if(r <= 55) k = (64 - 8 - 1) - r;  // 需要填充k个0（单位：字节）
    else k = (128 - 8 - 1) - r;     // k1 = 64 - r - 1, k2 = k - k1 - 8
    *padded_size = message_size + 1 + k + 8;
    padded = (uint8_t *)malloc(*padded_size); // 最后多加1是避免最后一个bit不是空字符存在安全问题
    memset(padded, 0, message_size+1+k+8);       // 清零

    memcpy(padded, message, message_size);
    *(padded+message_size) = 0x80;
    for(int i = 1; i <= k; ++i) {
        memset(padded+(message_size+i), 0, 1);
    }
    //memset(padded+(message_size+k+1), 0, 4); // 前面四个字节为0
    //*(uint32_t *)(padded+(message_size+k+2)) = message_size;
    uint64_t val = (uint64_t)message_size * 8;

    for(int i=7; i>=0; i--){
        *(padded+(message_size + k + 1 + i)) = (uint8_t)(val >> (56 - (i<<3)));
    }
    // *((padded+(message_size + k + 1))) = (uint32_t)(val>>32);
    // *((uint32_t*)(padded+(message_size + k + 1)) + 1) = (uint32_t)val;
    // memcpy(padded + (message_size + k + 1), &val, sizeof(val));
    return padded;
}

// 压缩函数，核心的哈希计算步骤，适配32位系统
void compress(uint32_t *V, uint8_t *block) {
    uint32_t W[68];
    uint32_t W_prime[64];
    // uint32_t *b = (uint32_t *)block;
    // 消息扩展，生成W和W'数组
    for (int i = 0; i < 16; i++) {
        W[i] = \
        ((uint32_t)block[4 * i] << 24) | \
        ((uint32_t)block[4 * i + 1] << 16) | \
        ((uint32_t)block[4 * i + 2] << 8) | \
        ((uint32_t)block[4 * i + 3]);
    }
    for (int i = 16; i < 68; i++) {
        W[i] = P1(W[i - 16] ^ W[i - 9] ^ ROTL(W[i - 3], 15)) ^ ROTL(W[i - 13], 7) ^ W[i - 6];
    }
    for (int i = 0; i < 64; i++) {
        W_prime[i] = W[i] ^ W[i + 4];
    }
    // 压缩计算
    uint32_t A = V[0];
    uint32_t B = V[1];
    uint32_t C = V[2];
    uint32_t D = V[3];
    uint32_t E = V[4];
    uint32_t F = V[5];
    uint32_t G = V[6];
    uint32_t H = V[7];

    uint32_t SS1, SS2, TT1, TT2;
    uint32_t T_j;
    for (int i = 0; i < 64; i++) {
        // if(i >= 0 && i <= 15){
        //     SS1 = ROTL(ROTL(A, 12) + E + ROTL(0x79CC4519, i), 7);
        // }else{
        // }
        T_j = (i <= 15) ? T_j_1 : T_j_2;
        SS1 = ROTL(ROTL(A, 12) + E + ROTL(T_j, i & 0x1F), 7);

        SS2 = SS1 ^ ROTL(A, 12);
        TT1 = FF(A, B, C, i) + D + SS2 + W_prime[i];
        TT2 = GG(E, F, G, i) + H + SS1 + W[i];
        D = C;
        C = ROTL(B, 9);
        B = A;
        A = TT1;
        H = G;
        G = ROTL(F, 19);
        F = E;
        E = P0(TT2);
    }
    V[0] ^= A;
    V[1] ^= B;
    V[2] ^= C;
    V[3] ^= D;
    V[4] ^= E;
    V[5] ^= F;
    V[6] ^= G;
    V[7] ^= H;
}

// 按照接口实现的SM3算法主函数
void sm3_hash(uint8_t digest[32], const uint8_t *message, size_t message_size) {
    // uint8_t padded_message[SM3_BLOCK_SIZE * 2];  // 足够容纳填充后的消息
    uint8_t *padded_message = NULL; 
    size_t padded_size;
    padded_message = padding(message, message_size, &padded_size);

    uint32_t V[8];
    memcpy(V, IV, sizeof(IV));

    for (size_t i = 0; i < padded_size / SM3_BLOCK_SIZE; i++) {
        compress(V, padded_message + (i << 6)); //填充后的消息按512bit (64字节)
    }

    for (int i = 0; i < 8; i++) {
        digest[4 * i] = (uint8_t)(V[i] >> 24);
        digest[4 * i + 1] = (uint8_t)(V[i] >> 16);
        digest[4 * i + 2] = (uint8_t)(V[i] >> 8);
        digest[4 * i + 3] = (uint8_t)(V[i]);
    }
}

